{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "916a1a13",
   "metadata": {},
   "source": [
    "##  Notebook 2: Filling RAG outputs For Evaluation\n",
    "\n",
    "In this notebook, we will use the example RAG pipeline to populate the RAG outputs: contexts (retrieved relevant documents) and answer (generated by RAG pipeline).\n",
    "\n",
    "The example RAG pipeline provided as part of this repository uses langchain to build a chatbot that references a custom knowledge base. \n",
    "\n",
    "If you want to learn more about how the example RAG works, please see [03_llama_index_simple.ipynb](../notebooks/03_llama_index_simple.ipynb).\n",
    "\n",
    "- **Steps 1-5**: Build the RAG pipeline.\n",
    "- **Step 6**: Build the Query Engine, exposing the Retriever and Generator outputs\n",
    "- **Step 7**: Fill the RAG outputs "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f337c81-5975-4708-b3c7-7600f5c62b06",
   "metadata": {},
   "source": [
    "#### Steps 1-5: Build the RAG pipeline\n",
    "\n",
    "#### Define the LLM\n",
    "Here we are using a local llm on triton and the address and gRPC port that the Triton is available on. \n",
    "\n",
    "***If you are using AI Playground (no local GPU) replace, the code in the cell below with the following: ***\n",
    "```\n",
    "import os\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings\n",
    "os.environ['NVAPI_KEY'] = \"REPLACE_WITH_YOUR_API_KEY\"\n",
    "\n",
    "nv_embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\")\n",
    "nv_document_embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\", model_type=\"passage\")\n",
    "nv_query_embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\", model_type=\"query\")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90478b29-38eb-4610-a8e6-2bcdae7308fc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! pip install langchain --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe89593",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_community.document_loaders import DirectoryLoader\n",
    "from langchain.document_loaders import TextLoader\n",
    "from langchain_community.document_loaders import UnstructuredHTMLLoader\n",
    "from langchain_community.document_loaders import UnstructuredPDFLoader\n",
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "from langchain.document_loaders import UnstructuredFileLoader, Docx2txtLoader\n",
    "from langchain.vectorstores import FAISS\n",
    "import pickle\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "import re\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "472e609f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "# from triton_trt_llm import TensorRTLLM\n",
    "# trtllm =TensorRTLLM(server_url=\"llm:8001\", model_name=\"ensemble\", tokens=300)\n",
    "\n",
    "DOCS_DIR = \"../ORAN_kb/\" #\n",
    "DOCS_DIR = \"../ORAN_kb/\"\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "############################################\n",
    "# Component #2 - Embedding Model and LLM\n",
    "############################################\n",
    "os.environ['NVIDIA_API_KEY'] = \"\"\n",
    "os.environ['NVAPI_KEY'] = \"\"\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings\n",
    "\n",
    "# make sure to export your NVIDIA AI Playground key as NVIDIA_API_KEY!\n",
    "llm = ChatNVIDIA(model=\"mixtral_8x7b\")\n",
    "nv_embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\")\n",
    "nv_document_embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\", model_type=\"passage\")\n",
    "nv_query_embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\", model_type=\"query\")\n",
    "\n",
    "model_name = \"intfloat/e5-large-v2\"\n",
    "model_kwargs = {\"device\": \"cuda\"}\n",
    "encode_kwargs = {\"normalize_embeddings\": True}\n",
    "e5_embeddings = HuggingFaceEmbeddings(\n",
    "    model_name=model_name,\n",
    "    model_kwargs=model_kwargs,\n",
    "    encode_kwargs=encode_kwargs,\n",
    ")\n",
    "\n",
    "from langchain.embeddings import HuggingFaceBgeEmbeddings\n",
    "model_name = \"facebook/dragon-plus-context-encoder\"\n",
    "model_kwargs = {'device': 'cuda'}\n",
    "encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity\n",
    "dragon_embeddings = HuggingFaceBgeEmbeddings(\n",
    "    model_name=model_name,\n",
    "    model_kwargs=model_kwargs,\n",
    "    encode_kwargs=encode_kwargs,\n",
    "    query_instruction=\"Represent this sentence for searching relevant passages: \"\n",
    ")\n",
    "dragon_embeddings.query_instruction = \"Represent this sentence for searching relevant passages: \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d7bb49e-8a34-4953-a9eb-2311d1b7cd2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(list(llm.available_models))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc40755",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('facebook/dragon-plus-query-encoder')\n",
    "query_encoder = AutoModel.from_pretrained('facebook/dragon-plus-query-encoder')\n",
    "context_encoder = AutoModel.from_pretrained('facebook/dragon-plus-context-encoder')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f8ff005",
   "metadata": {},
   "source": [
    "#### Create a Prompt Template\n",
    "\n",
    "A [**prompt template**](https://gpt-index.readthedocs.io/en/latest/core_modules/model_modules/prompts.html) is a common paradigm in LLM development.\n",
    "\n",
    "They are a pre-defined set of instructions provided to the LLM and guide the output produced by the model. They can contain few shot examples and guidance and are a quick way to engineer the responses from the LLM. Llama 2 accepts the [prompt format](https://huggingface.co/blog/llama2#how-to-prompt-llama-2) shown in `LLAMA_PROMPT_TEMPLATE`, which we manipulate to be constructed with:\n",
    "- The system prompt\n",
    "- The context\n",
    "- The user's question\n",
    "  \n",
    "Much like LangChain's abstraction of prompts, LlamaIndex has similar abstractions for you to create prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33378c44",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# import the relevant libraries\n",
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "LLAMA_PROMPT_TEMPLATE = (\n",
    " \"<s>[INST] <<SYS>>\"\n",
    " \"Use the following context to answer the user's question. If you don't know the answer, just say that you don't know, don't try to make up an answer.\"\n",
    " \"<</SYS>>\"\n",
    " \"<s>[INST] Context: {context} Question: {question} Only return the helpful answer below and nothing else. Helpful answer:[/INST]\"\n",
    ")\n",
    "\n",
    "LLAMA_PROMPT = PromptTemplate.from_template(LLAMA_PROMPT_TEMPLATE)\n",
    "llm_llama = ChatNVIDIA(\n",
    "    model=\"playground_llama2_70b\",\n",
    "    temperature=0.2,\n",
    "    max_tokens=300\n",
    ")\n",
    "\n",
    "llm_mixtral = ChatNVIDIA(model=\"mixtral_8x7b\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52d16c40",
   "metadata": {},
   "source": [
    "### Load Documents\n",
    "Follow the step number 1 [defined here](../notebooks/05_dataloader.ipynb) to upload the pdf's to Milvus server.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15cc3f9d",
   "metadata": {},
   "source": [
    "In this rest of this section, we will load and split the pdfs of ORAN pdfs. We will use the `RecursiveCharacterTextSplitter`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "349f7235",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to the vector store file\n",
    "vector_store_path = \"vectorstore.pkl\"\n",
    "DOCS_DIR = \"../ORAN_kb\"\n",
    "# Load raw documents from the directory\n",
    "text_loader_kwargs={'autodetect_encoding': True} #loader_kwargs=text_loader_kwargs\n",
    "raw_txts = DirectoryLoader(DOCS_DIR, glob=\"**/*.txt\", show_progress=True, loader_cls=TextLoader).load()\n",
    "raw_htmls = DirectoryLoader(DOCS_DIR, glob=\"**/*.html\", show_progress=True, loader_cls=UnstructuredHTMLLoader).load()\n",
    "raw_pdfs = DirectoryLoader(DOCS_DIR, glob=\"**/*.pdf\", show_progress=True, loader_cls=UnstructuredPDFLoader).load()\n",
    "raw_docs = DirectoryLoader(DOCS_DIR, glob=\"**/*.docx\", show_progress=True, loader_cls=Docx2txtLoader).load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5c9a665",
   "metadata": {},
   "source": [
    "#### Generate and Store Embeddings\n",
    "##### a) Generate Embeddings \n",
    "[Embeddings](https://python.langchain.com/docs/modules/data_connection/text_embedding/) for documents are created by vectorizing the document text; this vectorization captures the semantic meaning of the text. \n",
    "\n",
    "We will use [intfloat/e5-large-v2](https://huggingface.co/intfloat/e5-large-v2) or nvidia retriever for the embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b436a12f-4462-4fcb-b34a-7d29dcac531f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_line_break(text):\n",
    "    text = text.replace(\"\\n\", \" \").strip()\n",
    "    text = re.sub(\"\\.\\.+\", \"\", text)\n",
    "    text = re.sub(\" +\", \" \", text)\n",
    "    return text\n",
    "    \n",
    "def remove_two_points(text):\n",
    "    text = text.replace(\"..\",\"\")\n",
    "    return text\n",
    "    \n",
    "def remove_two_slashes(text):\n",
    "    text = text.replace(\"__\",\"\")\n",
    "    return text\n",
    "\n",
    "def just_letters(text):\n",
    "    return re.sub(r\"[^a-z]+\", \"\", text).strip()\n",
    "\n",
    "def remove_non_english_letters(text):\n",
    "    return re.sub(r\"[^\\x00-\\x7F]+\", \"\", text)\n",
    "\n",
    "def langchain_length_function(text):\n",
    "    return len(just_letters(remove_line_break(text)))\n",
    "\n",
    "def word_count(text):\n",
    "    text = remove_line_break(text)\n",
    "    text = just_letters(text)\n",
    "    tokenizer = RegexpTokenizer(r\"[A-Za-z]\\w+\")\n",
    "    tokenized_text = tokenizer.tokenize(text)\n",
    "    tokenized_text_len = len(tokenized_text)\n",
    "    return tokenized_text_len\n",
    "\n",
    "def truncate(text, max_word_count=1530):\n",
    "    tokenizer = RegexpTokenizer(r\"[A-Za-z]\\w+\")\n",
    "    tokenized_text = tokenizer.tokenize(text)\n",
    "    return \" \".join(tokenized_text[:max_word_count])\n",
    "\n",
    "def strip_non_ascii(string):\n",
    "    ''' Returns the string without non ASCII characters'''\n",
    "    stripped = (c for c in string if 0 < ord(c) < 127)\n",
    "    return ''.join(stripped)\n",
    "\n",
    "def generate_questions(chunk, triton_client):\n",
    "    triton_client.predict_streaming(chunk)\n",
    "    chat_history = \"\"\n",
    "    response_streaming_list = []\n",
    "    while True:\n",
    "        try:\n",
    "            response_streaming = triton_client.user_data._completed_requests.get(block=True)\n",
    "            \n",
    "        except Exception:\n",
    "            triton_client.close_streaming()\n",
    "            break\n",
    "    \n",
    "        if type(response_streaming) == InferenceServerException:\n",
    "            print(\"err\")\n",
    "            triton_client.close_streaming()\n",
    "            break\n",
    "    \n",
    "        if response_streaming is None:\n",
    "            triton_client.close_streaming()\n",
    "            break\n",
    "    \n",
    "        else:\n",
    "            response_streaming_list.append(response_streaming)\n",
    "            chat_history = triton_client.prepare_outputs(response_streaming_list)\n",
    "            yield chat_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa5cbba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check for existing vector store file\n",
    "vector_store_exists = os.path.exists(vector_store_path)\n",
    "vectorstore = None\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100, length_function=langchain_length_function, )\n",
    "\n",
    "all_documents=raw_pdfs+raw_docs\n",
    "\n",
    "# split all the docs\n",
    "docs = text_splitter.split_documents(all_documents)\n",
    "print(len(docs))\n",
    "print(docs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eed40917-c5b4-42b8-a81d-de5ef573c801",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = docs\n",
    "#remove short chuncks\n",
    "filtered_documents = [item for item in documents if len(item.page_content) >= 200]\n",
    "[(len(item.page_content),item.page_content) for item in documents]\n",
    "documents = filtered_documents\n",
    "pd.DataFrame([doc.metadata for doc in documents])['source'].unique()\n",
    "#remove line break\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_line_break(documents[i].page_content)\n",
    "#remove two points\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_two_points(documents[i].page_content)\n",
    "#remove non english characters points\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_two_slashes(documents[i].page_content)\n",
    "#remove two points\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_two_points(documents[i].page_content)\n",
    "[(len(item.page_content),item.page_content) for item in documents]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9bfcb20",
   "metadata": {},
   "source": [
    "##### b) Store Embeddings \n",
    "\n",
    "We will use the LlamaIndex module [`ServiceContext`](https://gpt-index.readthedocs.io/en/latest/core_modules/supporting_modules/service_context.html) to bundle commonly used resources during the indexing and querying stage. \n",
    "\n",
    "In this example, we bundle the build resources: the LLM, the embedding model, the node parser, and the prompt helper.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94c619cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = documents\n",
    "print(len(docs))\n",
    "vectorstore_e5 = FAISS.from_documents(docs, e5_embeddings)\n",
    "print(\"Done e5\")\n",
    "vectorstore_dragon = FAISS.from_documents(docs, dragon_embeddings)\n",
    "print(\"Done DRAGON\")\n",
    "# vectorstore_nvolve = FAISS.from_documents(docs, embedding=NVIDIAEmbeddings(model=\"nvolveqa_40k\"))\n",
    "# vectorstore_nvolve = FAISS.from_documents(docs, nv_embedder)\n",
    "# print(\"Done Nvovlve\")\n",
    "vector_store_path = \"vectorstore_e5.pkl\"\n",
    "with open(vector_store_path, \"wb\") as f:\n",
    "    pickle.dump(vectorstore_e5, f)\n",
    "vector_store_path = \"vectorstore_dragon.pkl\"\n",
    "with open(vector_store_path, \"wb\") as f:\n",
    "    pickle.dump(vectorstore_dragon, f)\n",
    "    # pickle.dump(vectorstore_nvolve, f)\n",
    "\n",
    "file = open(\"./vectorstore_e5.pkl\",'rb')\n",
    "vectorstore_e5 = pickle.load(file)\n",
    "file.close()\n",
    "file = open(\"./vectorstore_dragon.pkl\",'rb')\n",
    "vectorstore_dragon = pickle.load(file)\n",
    "file.close()\n",
    "\n",
    "print(len(vectorstore_e5.docstore._dict))\n",
    "print(len(vectorstore_dragon.docstore._dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d149cb-b1bc-408d-89e0-4be9852a1839",
   "metadata": {},
   "outputs": [],
   "source": [
    "#nvidia embeddings\n",
    "docs = documents\n",
    "vector_store_path = \"vectorstore_nv.pkl\"\n",
    "vectorstore_nvolve = FAISS.from_documents(docs, nv_embedder)\n",
    "print(len(vectorstore_nvolve.docstore._dict))\n",
    "with open(vector_store_path, \"wb\") as f:\n",
    "    pickle.dump(vectorstore_nvolve, f)\n",
    "\n",
    "# file = open(\"./vectorstore_nv.pkl\",'rb')\n",
    "# vectorstore_nvolve = pickle.load(file)\n",
    "# file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7096611-8fd4-4d10-9e66-9d69dceb1936",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(vectorstore_nvolve.docstore._dict))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26a13f9b",
   "metadata": {},
   "source": [
    "### Step 6: Build the Query Engine, exposing the Retriever and Generator outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a97b50c2-e721-4338-aea5-1387abf875ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "import os\n",
    "os.environ['NVIDIA_API_KEY'] = \"\"\n",
    "os.environ['NVAPI_KEY'] = \"\"\n",
    "os.environ['NGC_API_KEY'] = \"\"\n",
    "\n",
    "llm = ChatNVIDIA(model=\"playground_llama2_70b\")\n",
    "# print(llm.get_model_details)\n",
    "result = llm.invoke(\"Write a ballad about LangChain.\")\n",
    "print(result.content)\n",
    "\n",
    "# retriever = vectorstore_nvolve.as_retriever()\n",
    "# prompt = ChatPromptTemplate.from_messages(\n",
    "#     [\n",
    "#         (\n",
    "#             \"system\",\n",
    "#             \"Answer solely based on the following context:\\n<Documents>\\n{context}\\n</Documents>\",\n",
    "#         ),\n",
    "#         (\"user\", \"{question}\"),\n",
    "#     ]\n",
    "# )\n",
    "\n",
    "# chain = (\n",
    "#     {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "#     | prompt\n",
    "#     | llm\n",
    "#     | StrOutputParser()\n",
    "# )\n",
    "\n",
    "# chain.invoke(\"where did harrison work?\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79980087",
   "metadata": {},
   "outputs": [],
   "source": [
    "############################################\n",
    "# Component #4 - LLM Response Generation and Chat\n",
    "############################################\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are a helpful AI assistant named Envie. You will reply to questions only based on the context that you are provided. If something is out of context, you will refrain from replying and politely decline to respond to the user.\"), (\"user\", \"{input}\")]\n",
    ")\n",
    "user_input = \"What are the restrictions on using the material from the O-RAN ALLIANCE e.V. specification?\"\n",
    "chain = prompt_template | llm | StrOutputParser()\n",
    "print(list(llm.available_models))\n",
    "print(llm.model)\n",
    "if user_input and vectorstore_e5!=None:\n",
    "    retriever = vectorstore_e5.as_retriever(search_kwargs={\"k\": 3, \"score_threshold\": 0.5})\n",
    "    docs = retriever.get_relevant_documents(user_input)\n",
    "    \n",
    "    context = \"\"\n",
    "    for doc in docs:\n",
    "        context += doc.page_content + \"\\n\\n\"\n",
    "    \n",
    "    augmented_user_input = \"Context: \" + context + \"\\n\\nQuestion: \" + user_input + \"\\n\"\n",
    "\n",
    "    full_response = \"\"\n",
    "\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response\n",
    "\n",
    "print(full_response)\n",
    "\n",
    "# from langchain.chains import RetrievalQA\n",
    "# import time\n",
    "\n",
    "# qa_chain = RetrievalQA.from_chain_type(\n",
    "#     llm,\n",
    "#     retriever=vectorstore_nvolve.as_retriever(search_kwargs={\"k\": 1, \"score_threshold\": 0.8}),\n",
    "#     chain_type_kwargs={\"prompt\": LLAMA_PROMPT}\n",
    "# )\n",
    "# start_time = time.time()\n",
    "# result = qa_chain({\"query\": user_input})\n",
    "# print(result)\n",
    "print(\"\\n\")\n",
    "print(context)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbf4bff6",
   "metadata": {},
   "source": [
    "### Step 7: Fill the RAG outputs \n",
    "\n",
    "Let's now query the RAG pipeline and fill the outputs `contexts` and `answer` on the evaluation JSON file.\n",
    "\n",
    "First, we need to load the previously generated dataset. So far, the RAG outputs fields are empty.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49e46b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the relevant libraries\n",
    "import json\n",
    "from IPython.display import JSON\n",
    "# load the evaluation data\n",
    "f = open('syn_data_oran.json')\n",
    "syn_data = json.load(f)\n",
    "print(len(syn_data))\n",
    "# show the first element\n",
    "JSON(syn_data[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df8a8e4e",
   "metadata": {},
   "source": [
    "Let now query the RAG pipeline and populate the `contexts` and `answer` fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54da631e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are a helpful and friendly intelligent AI assistant bot named ORAN Chatbot, deployed by the Artificial Intelligence Solutions Architecture and Engineering team at NVIDIA. The context given below will provide some documentation as well as ORAN specifications. Based on this context, answer the following question related to ORAN standards and specifications. If the question is not related to this, please refrain from answering.\"), (\"user\", \"{input}\")]\n",
    ")\n",
    "\n",
    "chain = prompt_template | llm | StrOutputParser()\n",
    "## Using e5 Retriever on Synthetic data\n",
    "for count,entry in enumerate(syn_data):\n",
    "    user_input = entry['question']\n",
    "    print(count, user_input)\n",
    "    retriever = vectorstore_e5.as_retriever(search_kwargs={\"k\": 5, \"score_threshold\": 0.8})\n",
    "    docs = retriever.get_relevant_documents(user_input)\n",
    "    context = \"\"\n",
    "    cons = []\n",
    "    for doc in docs:\n",
    "        context += doc.page_content + \"\\n\\n\"\n",
    "        cons.append(doc.page_content)\n",
    "#     print(context)\n",
    "    augmented_user_input = \"Context: \" + context + \"\\n\\nQuestion: \" + user_input + \"\\n\"\n",
    "    full_response = \"\"\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response   \n",
    "    # result = qa_chain({\"query\": user_input})\n",
    "    # entry[\"answer\"] = result['result']\n",
    "    entry[\"answer\"] = full_response\n",
    "#     print(entry[\"answer\"])\n",
    "    entry[\"contexts\"] = cons #[context]\n",
    "    \n",
    "print(syn_data[100])\n",
    "import json\n",
    "with open('eval_syn_QA_E5.json', 'w') as f:\n",
    "    json.dump(syn_data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6f425f4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Using dragon Retriever on Synthetic data\n",
    "for count,entry in enumerate(syn_data):\n",
    "    user_input = entry['question']\n",
    "    print(count, user_input)\n",
    "    retriever = vectorstore_dragon.as_retriever(search_kwargs={\"k\": 5, \"score_threshold\": 0.8})\n",
    "    docs = retriever.get_relevant_documents(user_input)\n",
    "    context = \"\"\n",
    "    cons = []\n",
    "    for doc in docs:\n",
    "        context += doc.page_content + \"\\n\\n\"\n",
    "        cons.append(doc.page_content)\n",
    "#     print(context)\n",
    "    augmented_user_input = \"Context: \" + context + \"\\n\\nQuestion: \" + user_input + \"\\n\"\n",
    "    full_response = \"\"\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response   \n",
    "    # result = qa_chain({\"query\": user_input})\n",
    "    # entry[\"answer\"] = result['result']\n",
    "    entry[\"answer\"] = full_response\n",
    "#     print(entry[\"answer\"])\n",
    "    entry[\"contexts\"] = cons #[context]\n",
    "    \n",
    "print(syn_data[100])\n",
    "import json\n",
    "with open('eval_syn_QA_Dragon.json', 'w') as f:\n",
    "    json.dump(syn_data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "357a8fd1-08d5-44c3-9741-68f744a06b19",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Using Nvolve Retriever on Synthetic data\n",
    "for count,entry in enumerate(syn_data):\n",
    "    user_input = entry['question']\n",
    "    print(count, user_input)\n",
    "    retriever = vectorstore_nvolve.as_retriever(search_kwargs={\"k\": 5, \"score_threshold\": 0.8})\n",
    "    docs = retriever.get_relevant_documents(user_input)\n",
    "    context = \"\"\n",
    "    cons = []\n",
    "    for doc in docs:\n",
    "        context += doc.page_content + \"\\n\\n\"\n",
    "        cons.append(doc.page_content)\n",
    "    print(\"\\n\", context)\n",
    "    augmented_user_input = \"Context: \" + context + \"\\n\\nQuestion: \" + user_input + \"\\n\"\n",
    "    full_response = \"\"\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response   \n",
    "    # result = qa_chain({\"query\": user_input})\n",
    "    # entry[\"answer\"] = result['result']\n",
    "    entry[\"answer\"] = full_response\n",
    "#     print(entry[\"answer\"])\n",
    "    entry[\"contexts\"] = cons #[context]\n",
    "    \n",
    "print(syn_data[100])\n",
    "import json\n",
    "with open('eval_syn_QA_Nvolve.json', 'w') as f:\n",
    "    json.dump(syn_data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80eef120-017e-4259-8016-d9678276918e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
